﻿//***************************************************************************************
// SobelFilter.cpp by Frank Luna (C) 2011 All Rights Reserved.
//***************************************************************************************

#include "sobel_filter.h"
#include <dx/core/dx_exception.h>

namespace ifire::dx {
SobelFilter::SobelFilter(
    ID3D12Device* device, UINT width, UINT height, DXGI_FORMAT format)
    : RenderTexture(device, width, height, format) {
  BuildResource();
}

CD3DX12_GPU_DESCRIPTOR_HANDLE SobelFilter::OutputSrv() { return mhGpuSrv; }

UINT SobelFilter::DescriptorCount() const { return 2; }

void SobelFilter::BuildDescriptors(CD3DX12_CPU_DESCRIPTOR_HANDLE hCpuDescriptor,
    CD3DX12_GPU_DESCRIPTOR_HANDLE hGpuDescriptor, UINT descriptorSize) {
  // Save references to the descriptors.
  mhCpuSrv = hCpuDescriptor;
  mhCpuUav = hCpuDescriptor.Offset(1, descriptorSize);
  mhGpuSrv = hGpuDescriptor;
  mhGpuUav = hGpuDescriptor.Offset(1, descriptorSize);

  BuildDescriptors();
}


void SobelFilter::Execute(ID3D12GraphicsCommandList* cmdList,
    ID3D12RootSignature* rootSig, ID3D12PipelineState* pso,
    CD3DX12_GPU_DESCRIPTOR_HANDLE input) {
  cmdList->SetComputeRootSignature(rootSig);
  cmdList->SetPipelineState(pso);

  cmdList->SetComputeRootDescriptorTable(0, input);
  cmdList->SetComputeRootDescriptorTable(2, mhGpuUav);

  auto to_access = CD3DX12_RESOURCE_BARRIER::Transition(mOutput.Get(),
      D3D12_RESOURCE_STATE_GENERIC_READ, D3D12_RESOURCE_STATE_UNORDERED_ACCESS);
  cmdList->ResourceBarrier(1, &to_access);

  // How many groups do we need to dispatch to cover image, where each
  // group covers 16x16 pixels.
  UINT numGroupsX = (UINT)ceilf(width_ / 16.0f);
  UINT numGroupsY = (UINT)ceilf(height_ / 16.0f);
  cmdList->Dispatch(numGroupsX, numGroupsY, 1);

  auto to_read = CD3DX12_RESOURCE_BARRIER::Transition(mOutput.Get(),
      D3D12_RESOURCE_STATE_UNORDERED_ACCESS, D3D12_RESOURCE_STATE_GENERIC_READ);
  cmdList->ResourceBarrier(1, &to_read);
}

void SobelFilter::BuildDescriptors() {
  D3D12_SHADER_RESOURCE_VIEW_DESC srvDesc = {};
  srvDesc.Shader4ComponentMapping = D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING;
  srvDesc.Format = format_;
  srvDesc.ViewDimension = D3D12_SRV_DIMENSION_TEXTURE2D;
  srvDesc.Texture2D.MostDetailedMip = 0;
  srvDesc.Texture2D.MipLevels = 1;

  D3D12_UNORDERED_ACCESS_VIEW_DESC uavDesc = {};

  uavDesc.Format = format_;
  uavDesc.ViewDimension = D3D12_UAV_DIMENSION_TEXTURE2D;
  uavDesc.Texture2D.MipSlice = 0;

  device_->CreateShaderResourceView(mOutput.Get(), &srvDesc, mhCpuSrv);
  device_->CreateUnorderedAccessView(
      mOutput.Get(), nullptr, &uavDesc, mhCpuUav);
}

void SobelFilter::BuildResource() {
  // Note, compressed formats cannot be used for UAV.  We get error like:
  // ERROR: ID3D11Device::CreateTexture2D: The format (0x4d, BC3_UNORM)
  // cannot be bound as an UnorderedAccessView, or cast to a format that
  // could be bound as an UnorderedAccessView.  Therefore this format
  // does not support D3D11_BIND_UNORDERED_ACCESS.

  D3D12_RESOURCE_DESC texDesc;
  ZeroMemory(&texDesc, sizeof(D3D12_RESOURCE_DESC));
  texDesc.Dimension = D3D12_RESOURCE_DIMENSION_TEXTURE2D;
  texDesc.Alignment = 0;
  texDesc.Width = width_;
  texDesc.Height = height_;
  texDesc.DepthOrArraySize = 1;
  texDesc.MipLevels = 1;
  texDesc.Format = format_;
  texDesc.SampleDesc.Count = 1;
  texDesc.SampleDesc.Quality = 0;
  texDesc.Layout = D3D12_TEXTURE_LAYOUT_UNKNOWN;
  texDesc.Flags = D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS;

  auto default_props = CD3DX12_HEAP_PROPERTIES(D3D12_HEAP_TYPE_DEFAULT);
  HASSERT(device_->CreateCommittedResource(&default_props, D3D12_HEAP_FLAG_NONE,
      &texDesc, D3D12_RESOURCE_STATE_GENERIC_READ, nullptr,
      IID_PPV_ARGS(&mOutput)));
}

} // namespace ifire